#!coding:utf-8
# 巡线算法，默认识别单线，使用hsv颜色空间(180,255,255)
from __future__ import print_function
import cv2
import numpy as np
import os
import yaml
import io
import json
import math

def analyze_quadrilateral(approx):
    """
    分析四边形类型
    """
    # 计算四边长度
    sides = []
    for i in range(4):
        side = np.linalg.norm(approx[i][0] - approx[(i+1)%4][0])
        sides.append(side)
    
    # 计算角度
    angles = []
    for i in range(4):
        p1 = approx[i][0]
        p2 = approx[(i+1)%4][0]
        p3 = approx[(i+2)%4][0]
        
        vector1 = p1 - p2
        vector2 = p3 - p2
        
        angle = math.degrees(math.atan2(vector2[1], vector2[0]) - math.atan2(vector1[1], vector1[0]))
        angle = angle % 360
        if angle > 180:
            angle = 360 - angle
        angles.append(angle)
    
    # 判断四边形类型
    if all(abs(side/sides[0] - 1) < 0.2 for side in sides):  # 所有边相等
        if all(abs(angle - 90) < 10 for angle in angles):  # 所有角为90度
            shape_name = "square"
        else:
            shape_name = "rhombus"
    elif abs(1 - sides[2]/sides[0]) < 0.2 and abs(1 - sides[3]/sides[1]) < 0.2:
        if all(abs(angle - 90) < 10 for angle in angles):  # 对边相等且所有角为90度
            shape_name = "rectangle"
        else:
            shape_name = "polygon"
    else:
        shape_name = "polygon"

    # details = f"边长: {sides[0]:.1f}, {sides[1]:.1f}, {sides[2]:.1f}, {sides[3]:.1f} | "
    # details += f"角度: {angles[0]:.1f}°, {angles[1]:.1f}°, {angles[2]:.1f}°, {angles[3]:.1f}°"
    
    return shape_name

class LineDetector:
    """
    LineDetector 类, 用来检测颜色
    Attributes:
    conf_file: str 配置文件
    min_cnt_area: int 最小轮廓数量,低于该值不会检测
    colors: {name:[(min_h,min_s,min_v),(max_h,max_s,max_v)]} hsv颜色阈值
    size: [width,height] 缩放尺寸
    roi: [x1,y1,x2,y2] 识别区域
    """

    def __init__(self):
        self.conf_file = os.path.dirname(
            os.path.abspath(__file__)) + "/default.yaml"
        stream = io.open(self.conf_file, 'r', encoding='utf-8')
        self.config = yaml.safe_load(stream)
        # self.min_cnt_area = self.config[u'最小轮廓']
        # self.colors = self.config[u'颜色区间']
        # 预定义颜色范围 (HSV格式)
        self.color_names = self.config[u'颜色名称']
        self.color_ranges = {
            'red': [
                (np.array([0, 100, 100]), np.array([10, 255, 255])),      # 红色低范围
                (np.array([160, 100, 100]), np.array([180, 255, 255]))    # 红色高范围
            ],
            'orange': [(np.array([11, 100, 100]), np.array([20, 255, 255]))],
            'yellow': [(np.array([21, 100, 100]), np.array([30, 255, 255]))],
            'green': [(np.array([31, 100, 100]), np.array([85, 255, 255]))],
            'cyan': [(np.array([86, 100, 100]), np.array([100, 255, 255]))],
            'blue': [(np.array([101, 100, 100]), np.array([120, 255, 255]))],
            'purple': [(np.array([121, 100, 100]), np.array([150, 255, 255]))],
            'magenta': [(np.array([151, 50, 100]), np.array([159, 255, 255]))],
            'black': [(np.array([0, 0, 0]), np.array([180, 255, 50]))],
            'gray': [(np.array([0, 0, 50]), np.array([180, 43, 200]))],
            'white': [(np.array([0, 0, 200]), np.array([180, 30, 255]))],
        }
        self.colors = {}
        for name in self.color_names:
            self.colors[name] = [ {"min": color[0], "max": color[1]} for color in self.color_ranges[self.color_names[name]] ]
        self.size = self.config[u'缩放尺寸']
        self.roi = self.config[u'识别区域']
        # print(self.colors)

        # 形状检测参数
        self.min_area = 200  # 最小轮廓面积
        self.max_area = 86400  # 最大轮廓面积 half of 480*360
        
        # 图像处理参数
        self.blur_kernel = 5
        self.dilate_iterations = 2
        self.erode_iterations = 1
        
        # 形状近似精度
        self.epsilon_factor = 0.02

    def detect_hsv(self, image, color_hsv):
        """
        detect_hsv 函数, 检测颜色
        Keyword arguments::
        image: image 输入opencv图像
        color_hsv: [(min_h,min_s,min_v),(max_h,max_s,max_v)] hsv颜色阈值
        Returns:
        cnt: 检测轮廓
        image: image 包含轮廓的图像
        """
        # image = self.cv_image
        # x1, x2, y1, y2 = self.roi[0], self.roi[1], self.roi[2], self.roi[3]
        # image = image[y1:y2,x1:x2]
        # image = self.rector.rect(image)
        # image = cv2.resize(image,(self.size[0],self.size[1]))
        # image = cv2.flip(image,-1)
        # cnt_r = detect_cnt(image,[])
        self.cnt = self.detect_cnt(image, color_hsv)
        # if cnt is not None:
        #     cv2.drawContours(image, [self.cnt], -1, (0, 255, 255), thickness=2)
        return self.toLineDetections(self.cnt), image
    def detect_color(self,image,color):
        # print(self.colors,color)
        if image is None:
            return [0,0,0,0,0],None
        color_hsv = self.colors[color]
        return self.detect_hsv(image,color_hsv)
    def detect_cnt(self, image, color_hsv):
        """
        detect_cnt 函数, 检测轮廓
        Keyword arguments::
        image: image 输入opencv图像
                color_hsv: [(min_h,min_s,min_v),(max_h,max_s,max_v)] hsv颜色阈值
                Returns:
                cnt: 检测轮廓
                """
        max_cnt = None
        mask = None
        # ret,thresh1 = cv2.threshold(image,127,255,cv2.THRESH_BINARY)
        # ret,thresh2 = cv2.threshold(image,127,255,cv2.THRESH_BINARY_INV)
        # return findMaxCnt(ret)

        # change to hsv model
        hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
        # get mask
        # print(len(color_hsv))
        if len(color_hsv) == 1:
            color = color_hsv[0]
            mask = cv2.inRange(hsv, np.array(
                color['min']), np.array(color['max']))
        elif len(color_hsv) >= 2:
            for color in color_hsv:
                maski = cv2.inRange(hsv, np.array(
                    color['min']), np.array(color['max']))
                if mask is None:
                    mask = maski
                else:
                    mask = cv2.bitwise_or(mask, maski)
        else:
            return max_cnt
        return self.findMaxCnt(mask)

    def findMaxCnt(self, mask):
        """
        findMaxCnt 函数, 查找最大轮廓
        Keyword arguments::
        mask: image 掩码图像
        Returns:
        max_cnt: 检测到的最大的轮廓
        """
        if mask is None:
            return None
        # 腐蚀操作
        # mask = cv2.erode(mask, None, iterations=2)
        # 膨胀操作，其实先腐蚀再膨胀的效果是开运算，去除噪点
        # mask = cv2.dilate(mask, None, iterations=2)
        # _,contours,hierarchy = cv2.findContours(mask,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)
        res = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
        if len(res) == 2:
            contours, hierarchy = res
        elif len(res) == 3:
            _, contours, hierarchy = res
        max_cnt = None
        max_area = 0
        cnts = []
        for cnt in contours:
            area = cv2.contourArea(cnt)
            if area > self.min_area and area < self.max_area and area > max_area:
                max_cnt = cnt
        return max_cnt
    def toLineDetections(self,cnt):
        if cnt is not None:
            center,wh,angle = cv2.minAreaRect(cnt)
            return [center[0],center[1],wh[0],wh[1],angle]
        return [0,0,0,0,0]
    def crop(self,image,x1,y1,x2,y2):
        if image is None:
            return None
        if y1 < y2 and x1 < x2:
            return image[y1:y2,x1:x2]
        return image

# class ColorShapeDetector:
#     def __init__(self):


        
    def detect_color_pixel(self, hsv_pixel):
        """检测单个像素的颜色"""
        h, s, v = hsv_pixel
        
        # 检查每种颜色范围
        for color_name, ranges in self.color_ranges.items():
            for lower, upper in ranges:
                if lower[0] <= h <= upper[0] and \
                   lower[1] <= s <= upper[1] and \
                   lower[2] <= v <= upper[2]:
                    return color_name
        
        return 'unknown'
    
    def get_color_mask(self, hsv_image, color_name):
        """获取指定颜色的掩码"""
        if color_name not in self.color_ranges:
            return None
        
        mask = np.zeros(hsv_image.shape[:2], dtype=np.uint8)
        ranges = self.color_ranges[color_name]
        
        for lower, upper in ranges:
            color_mask = cv2.inRange(hsv_image, lower, upper)
            mask = cv2.bitwise_or(mask, color_mask)
        
        return mask
    
    def preprocess_image(self, image):
        """图像预处理"""
        # 高斯模糊减少噪声
        blurred = cv2.GaussianBlur(image, (self.blur_kernel, self.blur_kernel), 0)
        
        # 转换为HSV颜色空间
        hsv = cv2.cvtColor(blurred, cv2.COLOR_BGR2HSV)
        
        return hsv, blurred
    
    def morphological_operations(self, mask):
        """形态学操作"""
        kernel = np.ones((5, 5), np.uint8)
        
        # 膨胀操作连接相邻区域
        dilated = cv2.dilate(mask, kernel, iterations=self.dilate_iterations)
        
        # 腐蚀操作去除噪声
        eroded = cv2.erode(dilated, kernel, iterations=self.erode_iterations)
        
        return eroded
    
    def detect_shapes(self, mask):
        """检测形状轮廓"""
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        shapes = []
        for contour in contours:
            # 计算轮廓面积
            area = cv2.contourArea(contour)
            
            # 过滤太小或太大的轮廓
            if area < self.min_area or area > self.max_area:
                continue
            
            center, wh, angle = cv2.minAreaRect(contour)
            if wh[0]*wh[1]/area > 2.5:
                continue

            # 计算轮廓周长
            perimeter = cv2.arcLength(contour, True)
            
            # 轮廓近似
            epsilon = self.epsilon_factor * perimeter
            approx = cv2.approxPolyDP(contour, epsilon, True)
            
            # 计算边界矩形
            x, y, w, h = cv2.boundingRect(contour)
            
            # 计算轮廓的中心点
            M = cv2.moments(contour)
            if M["m00"] != 0:
                cx = int(M["m10"] / M["m00"])
                cy = int(M["m01"] / M["m00"])
            else:
                cx, cy = x + w//2, y + h//2
            
            shapes.append({
                'contour': contour,
                'approx': approx,
                'area': area,
                'perimeter': perimeter,
                'bounding_rect': (x, y, w, h),
                'center': (cx, cy),
                'vertices': len(approx)
            })
        
        return shapes
    
    def classify_shape(self, shape_info):
        """分类形状"""
        vertices = shape_info['vertices']
        area = shape_info['area']
        perimeter = shape_info['perimeter']
        x, y, w, h = shape_info['bounding_rect']
        
        # 计算宽高比
        aspect_ratio = w / float(h) if h != 0 else 0
        
        # 计算圆形度
        circularity = 4 * np.pi * area / (perimeter * perimeter) if perimeter != 0 else 0
        
        # 根据顶点数分类
        if vertices == 3:
            return "triangle"
        elif vertices == 4:
            shape = analyze_quadrilateral(shape_info['approx'])
            if shape == 'polygon':
                center, wh, angle = cv2.minAreaRect(shape_info['contour'])
                if abs(wh[0]*wh[1]/area - 2) < 0.3 :
                    shape = "triangle"
            return shape
            # 检查是否是正方形或矩形
            if 0.8 <= aspect_ratio <= 1.2:
                return "square"
            else:
                return "rectangle"
        # elif vertices == 5:
        #     return "pentagon"
        # elif vertices == 6:
        #     return "hexagon"
        elif vertices >= 5:
            # 对于更多顶点的形状，检查圆形度
            if vertices >= 7 and circularity > 0.85:
                return "circle"
            elif vertices >= 7 and circularity > 0.7:
                return "oval"
            elif vertices <=6:
                center, wh, angle = cv2.minAreaRect(shape_info['contour'])
                if abs(wh[0]*wh[1]/area - 2) < 0.3 :
                    return "triangle"
                else:
                    return "polygon"
            else:
                return "polygon"
                # return f"polygon_{vertices}"
        else:
            # 对于顶点很少的形状，检查是否是圆形
            if circularity > 0.85:
                return "circle"
            elif circularity > 0.7:
                return "oval"
            else:
                return "unknown"
    
    def detect_color_shape(self, image, specific_color=None, specific_shape=None):
        """检测图像中的色块形状"""
        # 预处理图像
        hsv, blurred = self.preprocess_image(image)
        result_image = image.copy()
        
        detected_objects = []
        
        # 如果指定了特定颜色，只检测该颜色
        colors_to_detect = [specific_color] if specific_color else self.color_ranges.keys()
        
        for color_name in colors_to_detect:
            # 获取颜色掩码
            color_mask = self.get_color_mask(hsv, color_name)
            if color_mask is None:
                continue
            
            # 形态学操作
            processed_mask = self.morphological_operations(color_mask)
            
            # 检测形状
            shapes = self.detect_shapes(processed_mask)
            
            for shape_info in shapes:
                # 分类形状
                shape_type = self.classify_shape(shape_info)
                
                if shape_type == 'polygon' and len(shape_info['approx']) > 8:
                    continue

                # 获取形状的主要颜色（从中心点采样）
                cx, cy = shape_info['center']
                if 0 <= cx < hsv.shape[1] and 0 <= cy < hsv.shape[0]:
                    pixel_color = hsv[cy, cx]
                    detected_color = self.detect_color_pixel(pixel_color)
                else:
                    detected_color = color_name
                
                # 添加到检测结果
                detection = {
                    'color': detected_color,
                    'shape': shape_type,
                    'vertices' : shape_info['vertices'],
                    'approx': shape_info['approx'].tolist(),
                    'contour': shape_info['contour'],
                    'center': shape_info['center'],
                    'area': shape_info['area'],
                    'bounding_rect': shape_info['bounding_rect']
                }

                if specific_shape is None or specific_shape == shape_type:
                    # 在图像上绘制结果
                    self.draw_detection(result_image, detection)
                    del detection['contour']
                    detected_objects.append(detection)
        
        return result_image, detected_objects
    
    def draw_detection(self, image, detection):
        """在图像上绘制检测结果"""
        approx = detection['approx']
        color_name = detection['color']
        shape_type = detection['shape']
        contour = detection['contour']
        center = detection['center']
        x, y, w, h = detection['bounding_rect']
        
        # 设置颜色映射
        color_map = {
            'red': (0, 0, 255),
            'orange': (0, 165, 255),
            'yellow': (0, 255, 255),
            'green': (0, 255, 0),
            'cyan': (255, 255, 0),
            'blue': (255, 0, 0),
            'purple': (226, 43, 138),
            'magenta': (255, 0, 255),
            'white': (255, 255, 255),
            'black': (0, 0, 0),
            'gray': (128, 128, 128),
            'unknown': (128, 128, 128),
        }
        
        color = color_map.get(color_name, (128, 128, 128))
        
        # 绘制轮廓
        cv2.drawContours(image, [contour], -1, color, 2)
        
        # 绘制边界矩形
        # cv2.rectangle(image, (x, y), (x + w, y + h), color, 2)
        
        # 标记顶点
        # for point in approx:
        #     cv2.circle(image, tuple(point[0]), 3, (255, 0, 0), -1)

        # 绘制中心点
        cv2.circle(image, center, 5, color, -1)
        
        # 添加标签
        label = f"{color_name} {shape_type}"
        cv2.putText(image, label, (x, y - 10), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
        
        # 添加面积信息
        area_text = f"Area: {detection['area']:.0f}"
        cv2.putText(image, area_text, (x, y + h + 20), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.4, color, 1)
    

# 使用示例
def test_color_shape_detect():
    detector = LineDetector()
    
    # 方式1: 处理图像
    image = cv2.imread('color_shape_test.jpeg')
    if image is not None:
        result_img, detections = detector.detect_color_shape(image)
        print(f"检测到 {len(detections)} 个色块:")
        for i, obj in enumerate(detections):
            print(f"{i+1}. {obj['color']} {obj['shape']} (面积: {obj['area']}, 中点: {obj['center']})")
        
        cv2.imshow('Result', result_img)
        cv2.waitKey(0)
        cv2.destroyAllWindows()
    


if __name__ == '__main__':
    """
    测试函数，实时检测黄色
    """
    import time
    detector = LineDetector()
    capture = cv2.VideoCapture(0)
    while True:
        # test_image = cv2.imread('./test.png')
        ret, cv_image = capture.read()
        if cv_image is None:
            print('None')
            continue
        start = time.time()
        cnt, image = detector.detect_color(cv_image, u'黄色')
        end = time.time()
        print("detect 1 frame in %.2f ms" %
              ( (end - start)*1000))
        # cnt, image = detector.detect_hsv(cv_image, detector.colors[u'黄线'])
        cv2.imshow("images", image)
        # cv2.imshow("images", np.hstack([image, output]))
        c = cv2.waitKey(2)
        if c == 27:
            break
    capture.release()
